package com.share.dgs.instrumentation;

import graphql.execution.ExecutionStepInfo;
import graphql.execution.instrumentation.InstrumentationState;
import graphql.schema.DataFetchingEnvironment;

import java.util.HashMap;
import java.util.Map;
import java.util.UUID;

/**
 * @author caifan
 * @created 2021/3/12
 * description:
 */
public class CustomInstrumentationState implements InstrumentationState {
    private Map<String, Object> stateMap = new HashMap<>();

    private String traceId;

    public CustomInstrumentationState(String traceId, Map<String, Object> stateMap) {
        this.traceId = traceId;
        this.stateMap = stateMap;
    }

    public CustomInstrumentationState recordTiming(DataFetchingEnvironment dfe) {
        ExecutionStepInfo executionStepInfo = dfe.getExecutionStepInfo();
        stateMap.put("parentType", executionStepInfo.getParent());
        stateMap.put("fieldName", executionStepInfo.getField());
        traceId = UUID.randomUUID().toString();
        return new CustomInstrumentationState(traceId, stateMap);
    }

    public Map<String, Object> snapMap() {
        Map<String,Object> map = new HashMap<>();
        map.put("state", stateMap);
        map.put("traceId", traceId);
        return map;
    }
}
